#!/bin/sh

e2e_test_install() {
  install_metric_server
  install_keda
  install_vertical_pod_autoscaler \
    --set-string updater.extraArgs.in-recommendation-bounds-eviction-lifetime-threshold=2s \
    --set-string recommender.extraArgs.recommendation-margin-fraction=1.5

  kubectl create namespace "$CLUSTER_NAMESPACE"
  create_or_replace_cluster "$CLUSTER_NAME" "$CLUSTER_NAMESPACE" 2 \
    --set-string instanceProfiles[0].name=size-s \
    --set-string instanceProfiles[0].cpu=500m \
    --set-string instanceProfiles[0].memory=2Gi \
    --set nonProductionOptions.disablePatroniResourceRequirements=false \
    --set nonProductionOptions.disableClusterResourceRequirements=true \
    --set-string 'configurations.postgresconfig.postgresql\.conf.max_connections=104' \
    --set-string 'configurations.poolingconfig.pgBouncer.pgbouncer\.ini.default_pool_size=100' \
    --set-string 'cluster.autoscaling.mode=none'

  deploy_psql_pod "$CLUSTER_NAMESPACE"

  wait_pods_running "$CLUSTER_NAMESPACE" 3
  wait_cluster "$CLUSTER_NAME" "$CLUSTER_NAMESPACE"

  PGPASSWORD="$(kubectl get secret -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" \
    --template '{{ (index .data "superuser-password") | base64decode }}')"
  kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
    -- bash -c "PGPASSWORD=$PGPASSWORD pgbench -i -h $CLUSTER_NAME -s 10"
}

e2e_test_uninstall() {
  k8s_async_cleanup_namespace "$CLUSTER_NAMESPACE"
}

e2e_test() {
  run_test "Benchmarking cluster with vertical autoscaling" benchmark_cluster_vertical_autoscaling
}

benchmark_cluster_vertical_autoscaling() {
  create_or_replace_cluster "$CLUSTER_NAME" "$CLUSTER_NAMESPACE" 2 \
    --set-string instanceProfiles[0].name=size-s \
    --set-string instanceProfiles[0].cpu=500m \
    --set-string instanceProfiles[0].memory=2Gi \
    --set nonProductionOptions.disablePatroniResourceRequirements=false \
    --set nonProductionOptions.disableClusterResourceRequirements=true \
    --set-string 'configurations.postgresconfig.postgresql\.conf.max_connections=104' \
    --set-string 'configurations.poolingconfig.pgBouncer.pgbouncer\.ini.default_pool_size=100' \
    --set-string 'cluster.autoscaling.mode=vertical' \
    --set-string 'cluster.autoscaling.minAllowed.patroni.cpu=500m' \
    --set-string 'cluster.autoscaling.maxAllowed.patroni.cpu=2'

  wait_until kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
    -- bash -c "PGPASSWORD=$PGPASSWORD psql -q -h $CLUSTER_NAME-replicas -c 'SELECT 1'"
  kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
    -- bash -c "PGPASSWORD=$PGPASSWORD pgbench -T '$((E2E_TIMEOUT / 2))' -C -c 20 -j 20 --max-tries=1000000 --progress=5 --log --log-prefix=/tmp/pgbench_log --failures-detailed --aggregate-interval='$((E2E_TIMEOUT / 20))' -h $CLUSTER_NAME" 2>&1 | tee "$LOG_PATH/bench" &
  BENCH_PID="$!"
  trap_kill "$BENCH_PID"

  LATEST_CPU_QUANTITY='0.5'
  while true
  do
    cat "$SPEC_PATH/abstract/quantity.jq" > "$LOG_PATH/get-pod-name-with-patroni-container-more-than-X-cpu.jq"
    echo '.items[]|select(.spec.containers|any(.name == "patroni" and (.resources.limits.cpu | quantity) > '"$LATEST_CPU_QUANTITY"'))|.metadata.name' \
      >> "$LOG_PATH/get-pod-name-with-patroni-container-more-than-X-cpu.jq"
    if wait_until eval 'kubectl get pod -n "$CLUSTER_NAMESPACE" \
      -l "app=StackGresCluster,stackgres.io/cluster-name=$CLUSTER_NAME" -o json \
      | jq -f "$LOG_PATH/get-pod-name-with-patroni-container-more-than-X-cpu.jq" \
      | wc -l | grep -qxF 2 \
      || ! kill -0 "$BENCH_PID"'
    then
      if ! kill -0 "$BENCH_PID" >/dev/null 2>&1 \
        && ! grep -q '^command terminated with exit code [1-9]' "$LOG_PATH/bench-autoscaled"
      then
        break
      fi
      cat "$SPEC_PATH/abstract/quantity.jq" > "$LOG_PATH/get-cpu-quantity-with-patroni-container.jq"
      echo '.items[]|.spec.containers[]|select(.name == "patroni")|.resources.limits.cpu | quantity' \
        >> "$LOG_PATH/get-cpu-quantity-with-patroni-container.jq"
      LATEST_CPU_QUANTITY="$(kubectl get pod -n "$CLUSTER_NAMESPACE" \
        -l "app=StackGresCluster,stackgres.io/cluster-name=$CLUSTER_NAME" -o json \
        | jq -f "$LOG_PATH/get-cpu-quantity-with-patroni-container.jq" -r \
        | head -n 1)"
      AUTOSCALE_CHANGED="$(date +%s)"
      echo "$AUTOSCALE_CHANGED $LATEST_CPU_QUANTITY" >> "$LOG_PATH/autoscale-changed"
      echo "success VerticalPodAutoscaler scaled vertically the primary of the cluster based on CPU to $LATEST_CPU_QUANTITY"
    else
      echo "FAIL: VerticalPodAutoscaler can not scale vertically the primary of the cluster based on CPU"
      kubectl top pod --containers -n "$CLUSTER_NAMESPACE" -l "app=StackGresCluster,stackgres.io/cluster-name=$CLUSTER_NAME,stackgres.io/cluster=true" || true
      kill "$BENCH_PID" || true
      kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
        -- bash -c 'ps -ef | grep "[p]gbench" | tr -s " " | cut -d " " -f 2 | while read PID; do kill "$PID"; done' || true
      return 1
    fi
    kill "$BENCH_PID" || true
    kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
      -- bash -c 'ps -ef | grep "[p]gbench" | tr -s " " | cut -d " " -f 2 | while read PID; do kill "$PID"; done' || true
    kubectl exec -n "$CLUSTER_NAMESPACE" psql -q -- bash -c 'cat /tmp/pgbench_log*' > "$LOG_PATH/pgbench-transactions"

    wait_until kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
      -- bash -c "PGPASSWORD=$PGPASSWORD psql -q -h $CLUSTER_NAME-replicas -c 'SELECT 1'"
    kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
      -- bash -c "PGPASSWORD=$PGPASSWORD pgbench -T '$((E2E_TIMEOUT / 2))' -C -c 20 -j 20 --max-tries=1000000 --progress=5 --log --log-prefix=/tmp/pgbench_log --failures-detailed --aggregate-interval='$((E2E_TIMEOUT / 20))' -h $CLUSTER_NAME" 2>&1 | tee "$LOG_PATH/bench-autoscaled" &
    BENCH_PID="$!"
    trap_kill "$BENCH_PID"
  done

  if wait_until eval '! kill -0 "$BENCH_PID"' \
    && ! grep -q '^command terminated with exit code [1-9]' "$LOG_PATH/bench-autoscaled"
  then
    echo "success Benchmark terminated succesfully"
  else
    echo "FAIL: Benchmark not terminated succesfully"
    kill "$BENCH_PID" || true
    kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
      -- bash -c 'ps -ef | grep "[p]gbench" | tr -s " " | cut -d " " -f 2 | while read PID; do kill "$PID"; done' || true
    return 1
  fi
  kill "$BENCH_PID" || true
  kubectl exec -n "$CLUSTER_NAMESPACE" psql -q\
    -- bash -c 'ps -ef | grep "[p]gbench" | tr -s " " | cut -d " " -f 2 | while read PID; do kill "$PID"; done' || true
  kubectl exec -n "$CLUSTER_NAMESPACE" psql -q -- bash -c 'cat /tmp/pgbench_log*' > "$LOG_PATH/bench-autoscaled-transactions"
  echo "Test results before autoscaling"
  echo
  cat "$LOG_PATH/bench"
  echo
  echo "Transactions before autoscaling"
  echo
  cat "$LOG_PATH/pgbench-transactions"
  echo
  echo "Test results before autoscaling"
  echo
  cat "$LOG_PATH/bench-autoscaled"
  echo
  echo "Transactions before autoscaling"
  echo
  cat "$LOG_PATH/bench-autoscaled-transactions"
  echo
  echo "Autoscaling completed at $AUTOSCALE_START"
}
